majit: stop inheriting a previous compile's target tokens - #1287

A token-minting compile inherited the previous compile's target tokens. Ten commits: the
measured fix, the deviation that made it possible, and the citations that hid it.

Rebased onto origin/main (81e8a2d); git range-diff reports the nine pre-rebase commits
unchanged.

The defect

compile.py:245 and :290 assign jitcell_token.target_tokens a fresh single-element
list, so a token-minting compile carries only its own labels. Only compile.py:341's
compile_retrace, which appends to the same token, accumulates. That is why
warmstate.py:191-196's invalidation filter covers the whole list: the list lives on the object
that was invalidated.

compile_trace_inner seeded the prior entry's front_target_tokens into the unroll optimizer
unconditionally. Instrumenting each compile's token list:

compile_loop (62 ops)  token=1  fronts=[837680, 837552]
compile_loop (32 ops)  token=3  fronts=[837680, 837552, 837664]

Token 3 is the loop compiled after the store invalidates the first one, and it carries the dead
loop's two labels ahead of its own.

The damage is not at the mint's own close — that is already refused by the ownership gate in
jump_to_existing_trace_impl, whose attach_jitcell_token_number is None at every mint. The
seeded tokens are republished as the new loop's front_target_tokens, and later bridge
compiles, handed that list by &mut, match them and close into the prior loop's body. The
backend registry confirms the executed edge: all three post-store bridges bake the invalidated
module's slot, which then collects ~30.5M entries against a few hundred into the replacement.

The measured fix (9fb952f)

Apply the warmstate.py:191-196 filter at the seed.

synth/attr_cache_invalidation — the type version_tag half of the same mechanism — moves by
the same deltas (1002 → 602, 5 → 3). Two fixtures moving together is what makes this the
mechanism rather than a fixture special case.

The target is independently confirmed by a second instrument on a tree that has main + #1284
and not this PR. PYRE_WASM_TRACE_ENTRY_CENSUS=1 over the fixture, counting module entries
rather than ops:

trace_id=1 key=2   30,566,421     <- the invalidated loop
trace_id=4 key=0          602     <- its replacement
total              71,321,693     = 7.00 entries per outer iteration

So the entries this PR removes are 43% of the fixture's total, measured without reference to the
op accounting that found them.

On the ratio, stated carefully

30cd766 removes max-wasm-ratio=6 because the ratio came down, not because the ceiling went
up — it was the last such annotation in the tree. But the supporting ratio figure is
local aarch64 darwin and indicative only: this repo's own history says a local arm64 run
cannot grade the wasm ratio gate, and execution-only numbers from CI put the fixture at 5.23x on
main, not at the 7.8x a local run reads. CI on this PR is what decides whether removing the
annotation is justified. If it says otherwise, the annotation comes back — the ceiling does
not go up.

CI has now answered, and it answers in favour of removal. In pyre/check.py (ubuntu-24.04):

On main the annotation was actively suppressing a reading above the 4x default; on this PR the
fixture needs no allowance and nothing complains, which is the shape "the ratio came down" was
supposed to produce. The 3.3x→1.4x pair is cross-run, so read it as consistent with the
annotation's removal rather than as a measured speedup.

The one red on this PR is pre-existing on main. FAIL cranelift raise_catch exec 0.50s > pypy 0.18s ratio 2.9x > gate 2.5x, then cranelift 1 failed, 437 passed. main at 0d98226
fails the identical fixture, gate and backend at 3.3x, and main's last four completed runs
all failed. macOS in the same PR run passed 438/438 with raise_catch at 2.1x. It is a
main-wide gate miss, not this PR's to close here. Two other jobs (sandbox build + e2e,
pyre/check.py (windows-latest)) died in "Set up job" on 429 Too Many Requests / 503
fetching Swatinem/rust-cache — GitHub infrastructure, before any of this diff ran.

Two further cautions against over-reading the numbers above:

• Entries removed is not a proxy for ratio removed. Across fixtures the two anti-correlate:
  int_loop has ~1 entry total and the best ratio in the set (0.81x), while
  if_else_jump_forward has ~87M entries/dynasm-second at 1.82x. This fixture is a ~536M
  outlier, which is what makes it the right thing to attack — it is not a conversion rate.
  −51.4% ops and −43% entries should not be read as the ratio falling by either figure.
• There is currently no clean control. The newest completed main run predates jit: expand % and // by a constant on a backend with no mul-high #1284, and
  PR CI tests the merge ref, so a PR arm carries main commits the baseline arm lacks. Until a
  post-jit: expand % and // by a constant on a backend with no mul-high #1284 main run completes, a CI delta on this PR is not cleanly attributable.

The deviation behind it (20ffefb, 1b303ab, 634da99)

The filter removes only the invalidated subset; a live foreign token was still inherited,
which upstream never does. The seeding itself has no upstream counterpart. These three
commits close that, and they are parity-only — no .jitstats counter moves on any backend.

• 20ffefb — the mint inherits nothing. The binding, not the seed argument, is emptied,
  because the same list also feeds the republication fallback. ensure_preamble_target_token
  inserts the preamble token into an empty list, as
  test_ensure_preamble_target_token_inserts_start_descr_first already pins, so the seeded list
  is [start_descr] rather than an absent label.
• 1b303ab — compile_retrace has two arms. compile.py:355-356 resolves the token with
  get_procedure_token(greenkey) and asserts it, so the arm that reuses a token keeps the
  accumulated targets — that is orthodox. The arm without a resumekey mints a fresh token at
  compile.py:1013, and ResumeFromInterpDescr.compile_and_attach assigns that token's
  target_tokens nothing at all, so it starts owning nothing. (The commit message said it gets
  the :245/:290 treatment; those are in compile_simple_loop / compile_loop and not on
  this route — corrected in 4e96b65, and the correction strengthens the change rather than
  weakening it.) On that arm the same list also feeds the loop that rebinds each prior token's
  original_jitcell_token_number to the new number, which is what could make a retired token
  pass the ownership gate.
• 634da99 — history.py:499-503 gives TargetToken no producer field; upstream keeps it
  on the optimizer and the reference runs builder→token (shortpreamble.py:454-457), with
  inline_short_preamble discriminating by sb.target_token is target_token
  (unroll.py:376-385). pyre parked it on the token, and seed_prior_target_tokens stripped it
  off every seeded token for exactly the hazard this PR is about. Moving it to the phase-2
  Optimizer retires the strip — but the discrimination it provided is not free, since
  jump_to_existing_trace_impl walks every candidate, so the identity test replaces it in the
  same commit.

Two things stated precisely, because they are easy to overstate

The identity test is not is. descr_identity compares descriptor allocations, and
TargetToken clones share one Arc, so it answers equal across a clone family where
unroll.py:379 answers False for a distinct object. It is sufficient because
finalize_short_preamble mints a fresh LoopTargetDescr per compile and no token carries a
producer any more. The comment in the tree says that rather than calling it a spelling of is.

The producer's GC walk moved; it was not dropped. walk_rd_consts_refs reached the producer
through compiled_loops, which holds only post-compile copies — but shortpreamble.rs records
that a replay op is rooted by short_preamble_jump, whose only walk is that arm. The in-flight
optimizer's slot address is published for the duration of a compile, following
compile_snapshot_root_slots. That address names a local of the unroll call, which returns
before the compile entry does, so it is withdrawn by its own guard bound after that local — not
by CompileSnapshotRootsGuard, which drops later and would leave a root walk reading a dropped
local.

Corrections to comments this touched

• seed_prior_target_tokens said unroll.py:298 was the only setter of
  short_preamble_producer. unroll.py:507 in import_state is a second one.
• The comment above attach_jitcell_token_number cited compile.py:797-811 as compile_retrace;
  that range is AbstractResumeGuardDescr.compile_and_attach.
• The retrace mint is orthodox, and the comment that said otherwise is now gone (de4f9b2).
  It cited compile.py:266 — which is in compile_loop — and claimed "RPython avoids this
  entirely by reusing the same loop_jitcell_token". compile.py:392-393 dispatches
  compile_and_attach on the resumekey's class, and the two implementations map onto pyre's
  two arms exactly: AbstractResumeGuardDescr (:797-811) attaches under
  resumekey_original_loop_token without minting; ResumeFromInterpDescr (:1006-1022)
  mints at :1013. The token compile.py:355-356 resolves is the optimization-time one —
  it carries the closing JUMP descr and the retrace budget, not the installation identity.
  propagate_original_jitcell_token then runs on both arms (:806, :1014), so the
  re-stamping loop was never a consequence of pyre minting. That also makes the loop over
  prior_front_target_tokens provably dead once 1b303ab binds it to Vec::new(), so it is
  deleted rather than rewritten. Nothing else in this PR rested on the false version.
• compile.py:245 / :290 are not on the retrace route (4e96b65). They are the only two
  assignments of target_tokens upstream — the third writer is the history.py:440 class
  default None — but :245 is inside compile_simple_loop (:216-250) and :290 inside
  compile_loop (:251-340), while compile_retrace starts at :341. Four sites cited them,
  or ranges containing them, to describe the retrace path, including the comment 1b303ab itself
  added. The conclusion is strengthened, not weakened: upstream's minted token on that arm
  does not get [start_descr], it gets nothing, so "seed nothing" was under-argued. The citation
  was still wrong and is fixed. Two sites citing the same lines on their own routes
  (pyjitpl.rs:6507 in compile_loop_body, :9827 in compile_simple_loop) are correct and
  left alone.
• Two claims were wrong on the pyre side, not the upstream side (4e96b65).
  JitCellToken::target_tokens' doc said the list is populated so has_compiled_loop reads what
  upstream's has_compiled_targets reads. has_compiled_loop is
  entry_procedure_token(gk).is_some(), and pyre's has_compiled_targets reads
  compiled_loops[gk].front_target_tokens — neither touches this list. Its one reader is
  first_target_token. That function's doc in turn implied pyjitpl.py:3007 closes onto the head
  descr, where upstream passes the JitCellToken and unroll.py:320-340 selects among
  target_tokens by virtual-state match. Re-reading the cited upstream line catches neither,
  because the false half is on the pyre side.
• has_compiled_targets was cited as pyjitpl.py:3898 at seven sites; it is at :3922-3923
  (:3898 is handler.__name__ = 'handler_' + name).

The audit those corrections prompted (857c0bc)

Both clusters above were found by stumbling over them, which is a bad way to learn that a
citation is wrong. A citation that quotes the upstream statement is machine-checkable, so
the tree was audited on that basis: 9373 citations → 2209 carrying a quote → 827 gradeable
(the quote anchored on a code identifier, not prose) → 760 correct, 67 candidates. Each of the
67 was judged individually against the vendored source; 30 turned out correct (usually a quote
naming a class or function while the cite points inside its body), 36 were stale, and 1 named
the wrong file (optimizer.py:317 for what is unroll.py:317). All 37 are fixed here —
21 files, 37 insertions and 37 deletions, comments only.

By upstream file: history.py 14, blackhole.py 9, pyjitpl.py 6, compile.py 2,
rewrite.py 2, unroll.py 2, warmstate.py 1, plus the filename fix.

No mechanical pass is safe on this, which is why each site was verified. Much of the rot in
pyjitpl.py is a uniform +24 shift, but the stale and the correct interleave:
dont_trace_here is cited as :2822 (stale) at one site and :2846 (correct) at another, in
the same tree — the citations were written at different times against different vendored
revisions, so a sed over any line band would corrupt the correct ones. The quote is what
discriminates, which is also an argument for quoting the statement whenever citing upstream.

This is comment-only and unrelated to the target-token subject; it is folded in rather than
split out because the audit was a direct consequence of the two clusters above.

Refuted along the way, recorded so they are not re-proposed

• The close gate and the JUMP descr read two sources. Landed in wasm: pass a failing guard's fail args to its bridge as call parameters #1274; the entry census came
  back byte-identical.
• The prior_front_target_tokens fallback in the InvalidLoop path. An eprintln in both
  arms never fired.
• "There is no zombie." Read off an undecomposed counter that is a sum over modules.
  Retracted.
• "The mint's own close resolves into the invalidated loop." The ownership gate already
  refuses it; the channel is republication into later bridge compiles.

Still open, deliberately not touched here

compile_entry_bridge's inheritance — measured, and not this mechanism. It clones the
retired loop's front_target_tokens on the replace path. Rather than argue about it, a
PYRE_PROBE_EB counter at the site was run on both invalidation fixtures under dynasm release:
zero hits on both, with stdout still correct. So it is not implicated in what the −51.4%
addressed. Zero hits on two fixtures is not "dead in general", only "not this mechanism", and
CarriedFields' doc already names it as the one of five replace paths that legitimately
inherits.

Where pyre still records what upstream leaves empty. On the minting arm, pyre mirrors the
fresh target tokens onto JitCellToken.target_tokens; upstream's ResumeFromInterpDescr arm
never assigns that list, and history.py:440 leaves it None, so
has_compiled_targets(ptoken) is False for the token attach_procedure_to_interp installs.
This is not the "trace owns no LABEL" case it resembles — compile.py:382 puts [label_op]
in the retrace's operations, so upstream has a LABEL there and propagates through it; it simply
does not mirror. Closing the gap is not a deletion, because the two trees read the list through
different objects: upstream's JUMP carries the JitCellToken and unroll.py:320-325 walks
target_tokens at optimization time, whereas pyre resolves first_target_token() at record
time (compile_trace) and cancels the bridge when it is absent. Dropping the record here would
change which compiles happen, so it is a separate, measured change — not a comment fix, and out
of scope for this PR.

That question is now settled, against moving. Upstream's cell-token descr is a placeholder the
optimizer always consumes, in one of two ways: unroll.py:196-199 takes jump_to_preamble when
the list holds one entry, and :238-241 rewrites the descr to cell_token.target_tokens[0] —
element zero, unconditionally, which is exactly what first_target_token answers; otherwise
:320-359 virtual-state matches and rewrites to the token it picked. Both consumers exist
here. jump_to_existing_trace_impl iterates every candidate and its unroll.py:357-359 arm
re-points the JUMP's descr at whichever token matched, so binding early does not bypass the
ladder — only the preamble arm keeps what was recorded. The equivalence holds because recording
and optimizing are one synchronous sequence on the single JIT thread and optimize_bridge mints
no targets.

Adopting the literal upstream shape would not pay: it works upstream because
token.target_tokens holds the full TargetTokens the ladder consumes, and here it cannot —
JitCellToken lives in majit-backend while TargetToken-with-VirtualState lives in
majit-metainterp, which depends on it, the reverse of history.py owning both. The token can
only carry the descr projection, so a JitCellToken-descr'd JUMP would be traded back for the same
side-table list the optimizer already receives as a parameter. That is the RPython↔Rust layering
gap, and it is now cited at the site rather than left implicit.

What the exercise did surface is a real gap, and this PR closes it: nothing enforced that the
descr list on the token and the value list in compiled_loops agree. They are two projections
of one thing written by separate statements, and upstream cannot drift because it holds one list
of real TargetTokens. A debug_assert now checks that the resolved head equals
front_target_tokens[0] under descr_identity. cargo test is a debug build, so the workspace
suite exercises it.

Nothing above should be read as closing the rest: the −51.4% is a measured win on the mint
channel, not proof the others are absent.

Verification

Re-run on the tip after each commit, and once more on the rebased tip 857c0bc:

• cargo fmt --check — clean
• cargo test --workspace — green, 8073 tests
• python3 pyre/check.py --backend wasm --synthetic-only — 417/417
• python3 pyre/check.py --backend dynasm --synthetic-only — 421/421
• cargo test -p pyre-jit --features dynasm --test gc_stress — 34/34 (not re-run on
  857c0bc, which is comment-only)
• scripts/extract-llbc.py — clean, so pyre-jit.ullbc is not stale
• no .jitstats baseline re-recorded for any of the parity commits

Summary by CodeRabbit

• Performance

  • Improved JIT short-preamble handling during loop compilation and replay.
  • Reduced unnecessary guard failures and bridge compilations in affected benchmarks.

• Bug Fixes

  • Improved handling of invalidated loops and retracing.
  • Preserved compiled loop state more reliably during compilation and recovery.

• Documentation

  • Corrected numerous internal documentation references and explanatory comments.
  • Updated benchmark notes to reflect revised WebAssembly performance results.